package org.activequant.math.algorithms;

import org.apache.log4j.Logger;

import flanagan.math.Matrix;

/**
 * Fits sequence of N-vectors {y_0, y_1, ... y_n}, {x_0, x_1, ... x_n} into a linear model:
 * <pre>
 * y = alpha + phi x
 * </pre>
 * where alpha is N-vector, phi is NxN matrix.
 * <p>
 * Result is the learned matrix phi, vector alpha, and noise correlation matrix sigma.
 * <b>History:</b><br>
 *  - [10.02.2008] Created (Mike Kroutikov)<br>
 *
 *  @author Mike Kroutikov
 */
public class EMAMatrixRegression {
	
	private final Logger log = Logger.getLogger(getClass());
	
	private final EMAMatrixAccumulator x;
	private final EMAMatrixAccumulator f;
	private final EMAMatrixAccumulator b;
	private final EMAMatrixAccumulator i;
	private final EMAMatrixAccumulator k;

	public EMAMatrixRegression(int nrows, int ncols, int period) {
		x = new EMAMatrixAccumulator(nrows, 1);
		x.setPeriod(period);
		
		f = new EMAMatrixAccumulator(nrows, 1);
		f.setPeriod(period);

		b = new EMAMatrixAccumulator(nrows, ncols);
		b.setPeriod(period);
		
		i = new EMAMatrixAccumulator(nrows, ncols);
		i.setPeriod(period);
		
		k = new EMAMatrixAccumulator(nrows, ncols);
		k.setPeriod(period);
		
		phi   = new Matrix(nrows, ncols);
		alpha = new Matrix(nrows, 1);
		mu    = new Matrix(nrows, 1);
		sigma = new Matrix(nrows, ncols);
	}
	
	private static Matrix otimes(Matrix a, Matrix b) {
		if(a.getnumberOfColumns() != 1 || b.getnumberOfColumns() != 1) {
			throw new IllegalArgumentException("otimes takes Vectors, not Matrices!");
		}
		
		Matrix out = new Matrix(a.getnumberOfRows(), b.getnumberOfRows());
		for(int i = 0; i < a.getnumberOfRows(); i++) {
			for(int j = 0; j < b.getnumberOfRows(); j++) {
				out.setElement(i, j, a.getElement(i, 0) * b.getElement(j, 0));
			}
		}
		
		return out;
	}
	
	public void accumulate(Matrix yy, Matrix xx) {
		Matrix ff = Matrix.minus(yy, xx);

		Matrix xprev = x.getMeanValue().copy();
		x.accumulate(xx);

		Matrix fprev = f.getMeanValue().copy();
		f.accumulate(ff);
		
		b.accumulate(otimes(xx.minus(xprev), xx.minus(x.getMeanValue())));
		i.accumulate(otimes(ff.minus(fprev), ff.minus(f.getMeanValue())));
		k.accumulate(otimes(ff.minus(fprev), xx.minus(x.getMeanValue())));
		
		log.info("x=" + toString(x.getMeanValue()));
		log.info("f=" + toString(f.getMeanValue()));
		log.info("b=" + toString(b.getMeanValue()));
		log.info("k=" + toString(k.getMeanValue()));

		solved = false;
	}

	private Matrix phi;
	private Matrix alpha;
	private Matrix mu;
	private Matrix sigma;
	
	private boolean solved = false;
	
	private static final double EPSILON = 1.e-10;
	
	public boolean isReady() {
		double detB = b.getMeanValue().determinant();
		double detK = k.getMeanValue().determinant();
		detK = Math.abs(detK);
		
		return detB > EPSILON * detK && detK > EPSILON * detB;
	}
	
	private void solve() {
		Matrix revB = b.getMeanValue().inverse();
		phi = k.getMeanValue().times(revB);

		// alpha = f - phi * x
		alpha = f.getMeanValue().minus(phi.times(x.getMeanValue()));
		
		mu = x.getMeanValue().minus(b.getMeanValue().times(k.getMeanValue().inverse()).times(f.getMeanValue()));
	
		Matrix n = k.getMeanValue().times(revB);
		
		sigma = i.getMeanValue();
		
		sigma.minusEquals(n.times(k.getMeanValue()));
		
		// seetting solved at the end of this method. 
		solved=true;
	}
	
	public Matrix getPhi() {
		if(!solved) solve();
		return phi.plus(Matrix.identityMatrix(phi.getnumberOfRows()));
	}
	
	public Matrix getAlpha() {
		if(!solved) solve();
		return alpha;
	}

	public Matrix getMu() {
		if(!solved) solve();
		return mu;
	}

	public Matrix  getSigma() {
		if(!solved) solve();
		return sigma;
	}
	
	public Matrix getAutoCorrelator() { 
		return b.getMeanValue(); 
	}
	
	private static String toString(Matrix m) {
		String out = "Matrix(" + m.getnumberOfRows() + "x" + m.getnumberOfColumns() + "): ";
		for(int i = 0; i < m.getnumberOfRows(); i++) {
			out += "[";
			for(int j = 0; j < m.getnumberOfColumns(); j++) out += " " + m.getElement(i, j);
			out += "]";
		}
		
		return out;
	}

	public String toString() {
		if(!solved) solve();
		
		return "EMAMatrixRegression: period=" + x.getPeriod() 
			+ ", numSamples=" + x.getNumSamples()
			+ ", phi=" + toString(getPhi())
			+ ", alpha=" + toString(alpha)
			+ ", sigma=" + toString(sigma);
	}
}
